import json
import os
import random
import time
from abc import abstractmethod
from typing import List, Tuple, Iterable, Union
import h5py
import numpy as np

class TagsBase:
    def __init__(self) -> None:
        self.output = {"ReomteDataVersion": 0, "LocalDataVersion": 1, "EventData": []}
        self.is_full = True
    
    def set_chan(self, is_full) -> None:
        '''
        功能：设置事件通道类型
        参数：
            is_full：True表示多通道，False表示单通道，默认是多通道
        '''
        self.is_full = is_full
    
    def add_envent(self, event_name:str, 
                   start:Union[int, float], 
                   duration:Union[int, float], 
                   chan_num:int=None) -> None:
        '''
        功能：添加事件
        参数：
            event_name: 事件的名称
            start: 开始的时间
            duration: 结束的时间
            chan_num: 仅通道类型为单通道时需要，表示通道位置，从0开始
        '''
        item_id = str(int(time.time() * 100000000)+random.randint(1, 100000))
        item = {
            'Id': item_id,
            'EventName': event_name,
            'IsFull': self.is_full,
            'Start': start,
            'Duration': duration,
            'bgColorStr': '01ffffff'
        }
        if self.is_full == False:
            assert chan_num is None
            item['ChanNum'] = chan_num
        self.output['EventData'].append(item)
    
    def to_tags(self, out_folder=None, out_file=None, out_path=None) -> None:
        '''
        功能：输出tags文件
        参数：
            out_folder：输出文件夹
            out_file：输出文件名
            out_path：输出路径，如果有，优先使用
        '''
        json_str = json.dumps(self.output, ensure_ascii=False, indent=2)
        if out_path is None:
            out_path = os.path.join(out_folder, out_file)
        with open(out_path, 'w') as f:
            f.write(json_str)
        self.output = {"ReomteDataVersion": 0, "LocalDataVersion": 1, "EventData": []}
    
    @abstractmethod
    def collect_event(self):
        pass

    @abstractmethod
    def auto_to_tags(self):
        pass

class ToTags(TagsBase):
    def __init__(self, window:int, label_map:dict) -> None:
        '''
        参数：
            window：窗口大小
            label_map：映射标签
        '''
        self.window = window
        self.label_map = label_map
        self.SRR = 1
        super().__init__()
    
    def set_SRR(self, SRR:float) -> None:
        '''
        功能：动态设置SRR
        参数：
            SRR：旧采样率/新采样率，默认是1
        '''
        self.SRR = SRR

    def collect_event(self, events:Iterable, starts:Iterable, converted:bool, 
                      durations:Iterable=None) -> None:
        '''
        功能：收集事件
        参数：
            events：元素是事件编码
            starts：元素是开始时间
            converted：开始时间是否有转换
            durations：元素是结束时间
        '''
        logic = (converted, durations is None)
        for i, (event, start) in enumerate(zip(events, starts)):
            if event not in self.label_map:
                continue
            if logic == (True, True):
                duration = start + self.window * self.SRR
            elif logic == (False, True):
                duration = (start + self.window) * self.SRR
                start *= self.SRR
            elif logic == (False, False):
                duration = durations[i]
                duration *= self.SRR
                start *= self.SRR
            self.add_envent(self.label_map[event], int(start), int(duration))
    
    def auto_to_tags(self):
        pass

class H5ToTags(ToTags):
    def load_h5(self, h5_path) -> Tuple[np.ndarray, np.ndarray]:
        '''
        功能：读取h5格式文件数据
        参数：
            h5_path：h5文件路径
        返回：
            data：data字段数据
            label：label字段数据
        '''
        with h5py.File(h5_path) as f:
            data = f['data'][:]
            label = f['label'][:]
        return data, label
    
    def auto_to_tags(self, label:Iterable[Tuple[int, bytes]], out_folder, converted:bool=False) -> None:
        '''
        功能：自动生成tags文件
        参数：
            label：label字段数据
            out_folder：输出文件夹
            converted：开始时间是否有转换，默认是没有转换
        '''
        ids:List[str] = [i[1].decode('utf-8') for i in label]
        all_events = [i[0] for i in label]
        fids, all_starts = [], []
        for _id in ids:
            fid, start = _id.rsplit('_', 1)
            fids.append(fid)
            all_starts.append(start)
        end = fids[0]
        events, starts = [], []
        for fid, event, start in zip(fids, all_events, all_starts):
            if end != fid:
                self.collect_event(events, starts, converted=converted)
                self.to_tags(out_folder, end+'.tags')
                events, starts = [], []
                end = fid
            starts.append(float(start))
            events.append(event)
        # 最后一个文件
        self.collect_event(events, starts, converted=converted)
        self.to_tags(out_folder, fid+'.tags')

class CustomFileToTags(ToTags):
    def load_txt(self, txt_path):
        return np.loadtxt(txt_path, dtype=np.int64)
    
    def auto_to_tags(self, label: Iterable[Iterable[int]], out_folder, out_name, converted: bool = False) -> None:
        self.collect_event(label[:, 0], label[:, 1], label[:, 2], converted=converted)
        self.to_tags(out_folder=out_folder, out_file=out_name+'.tags')

if __name__ == '__main__':
    txt_path = './329-data8.txt'
    ftt = CustomFileToTags(100, {1:'尖慢波'})
    ftt.set_SRR(2.56)
    label = ftt.load_txt(txt_path)
    out_name = os.path.basename(txt_path).split('.')[0]
    ftt.auto_to_tags(label, '../data/temp', out_name, converted=False)
